# trajs loss
"""
---This file contains the implementation of the motion loss function.

---The motion loss function is used to evaluate the motion of the predicted trajectory
against the ground truth trajectory.

---The motion loss function is defined as the mean squared error between the predicted
trajectory and the ground truth trajectory.

---The motion loss function is implemented as a PyTorch module.

obstacle loss for box 
map loss for line
comfort_loss
diversity_loss
distance loss 
similarity loss for drive trajs
occ loss
"""
from ..builder import build_loss
from mmdet.models.builder import LOSSES
import torch
from mmdet3d.core.bbox.iou_calculators.iou3d_calculator import bbox_overlaps_3d
def get_bounding_box(box):
    x, y, z, l, w, h, heading, _, _, _, _ = box
    # 计算障碍物的4个底部角点
    corners = torch.tensor([
        [-l/2, -w/2, 0],  # 左后
        [ l/2, -w/2, 0],  # 右后
        [ l/2,  w/2, 0],  # 右前
        [-l/2,  w/2, 0],  # 左前
    ], dtype=torch.float32)
    # 旋转矩阵
    rotation_matrix = torch.tensor([
        [torch.cos(heading), -torch.sin(heading), 0],
        [torch.sin(heading),  torch.cos(heading), 0],
        [0, 0, 1]
    ], dtype=torch.float32)
    
    # 旋转并平移
    corners = corners @ rotation_matrix.T + torch.tensor([x, y, z], dtype=torch.float32)
    return corners
def IOU_bev(box1: torch.Tensor, box2:torch.Tensor):
    """计算两个边界框的BEV IOU
     box1: N1 x 7
     box2: N2 x 7           
    """
    # todo
    dist = torch.norm(box1[:,:2].unsqueeze(0) - box2[:,:2].unsqueeze(1),dim=-1)
    return dist

def box_loss(pred_trajs, boxes, eps=0.5):
    N, P, _ = pred_trajs.shape
    N1, T, _ = boxes.shape
    assert P == T, "The number of predicted points should be equal to the number of time steps in boxes"
    # 计算每个时刻的障碍物边界框角点
    box_corners = torch.stack([get_bounding_box(box) for box in boxes.view(-1, 11)])  # (N1*T, 4, 3)
    box_corners_tensor = box_corners.view(N1, T, 4, 3)  # (N1, T, 4, 3)
    # 维度转换for loss
    box_corners_tensor = box_corners_tensor.transpose(0, 1)  # (T, N1, 4, 3)
    # 扩展维度以便于广播
    pred_trajs_expanded = pred_trajs.unsqueeze(2).unsqueeze(3)  # (N, T, 1, 1, 3)

    # 计算每个时刻的轨迹点与对应障碍物的距离
    distances = torch.norm(pred_trajs_expanded - box_corners_tensor.unsqueeze(0), dim=-1)  # (N, P, N1, 4)

    # 找到每个点到每个障碍物的最小距离
    min_distances = distances.min(dim=-1)[0]  # (N, P, N1)

    # 计算损失
    loss = ((min_distances < eps) * (eps - min_distances) ** 2).sum()

    return loss
		
def BST_IOU(box1,box2,t=16):
    # box1: (B,N1,S,T, 7+x) source
    # box2: (B,N2,1,T, 7+x) target
    B, N1, S, T, _ = box1.shape
    B1, N2, _, T, _ = box2.shape
    iou = torch.zeros((B,S,T,N1,N2), device=box1.device, dtype=box1.dtype)
    for b in range(B):
        for s in range(S):
            for t in range(T):
                iou[b,s,t] = bbox_overlaps_3d(box1[b,:,s,t], box2[b,:,0,t])
    return iou
def Bbox_loss(pred_trajs, gt_boxes, eps=0.5):
    #
    B, N1, S, T, _ = pred_trajs.shape
    iou = BST_IOU(pred_trajs, gt_boxes)
    loss = iou.sum() / (B*S*N1)

    return loss
def point_to_segment_distance(points, segments):
    """计算点到多个线段的最小距离"""
    points = points.unsqueeze(1)  # (P, 1, 2)
    segment_starts = segments[:, 0, :]  # (S, 2)
    segment_ends = segments[:, 1, :]  # (S, 2)

    # 计算线段向量和点到线段起点的向量
    line_vecs = segment_ends - segment_starts  # (S, 2)
    line_vecs_sq = (line_vecs ** 2).sum(dim=1, keepdim=True)  # (S, 1)

    # 计算点到线段起点的向量
    point_vecs = points - segment_starts.unsqueeze(0)  # (P, S, 2)

    # 点到线段的投影比例
    t = (point_vecs.unsqueeze(2) @ line_vecs.unsqueeze(2)).squeeze(-1) / line_vecs_sq  # (P, S, 1)

    # 限制在[0, 1]之间
    t = torch.clamp(t, 0, 1)

    # 计算最近点
    nearest_points = segment_starts.unsqueeze(0) + t * line_vecs.unsqueeze(0)  # (P, S, 2)

    # 计算距离
    distances = torch.norm(points - nearest_points, dim=-1)  # (N, P, S)
    return distances.min(dim=-1)[0]  # 返回每个点到所有线段的最小距离

def calculate_map_loss(pred_trajs, map_data):
    loss = 0.0
    safety_margin = 0.2  # 安全边界
    lane_safety_margin = 0.1  # 针对车道线的更小安全边界

    # 准备线段数据
    segments = {}
    for key in map_data.keys():
        if key in ["curbs", "lane_lines", "center_line", "crosswalks"]:
            segments[key] = torch.stack([map_data[key][:-1], map_data[key][1:]], dim=1)  # (S, 2, 2)

    # 计算与路沿的loss
    distances = point_to_segment_distance(pred_trajs[:, :, :2].reshape(-1, 2), segments["curbs"])
    loss += (distances[distances < safety_margin] ** 2).sum()

    # 计算与车道线的loss，增加惩罚
    lane_distances = point_to_segment_distance(pred_trajs[:, :, :2].reshape(-1, 2), segments["lane_lines"])
    loss += (lane_distances[lane_distances < lane_safety_margin] ** 2).sum() * 10  # 增加惩罚因子

    # 计算与中心线的loss
    distances = point_to_segment_distance(pred_trajs[:, :, :2].reshape(-1, 2), segments["center_line"])
    loss += (distances ** 2).sum()

    # 计算与斑马线的loss
    distances = point_to_segment_distance(pred_trajs[:, :, :2].reshape(-1, 2), segments["crosswalks"])
    loss += (distances[distances < safety_margin] ** 2).sum() * 3

    # 计算与红绿灯的loss
    for signal in map_data["traffic_signals"]:
        distances = torch.norm(pred_trajs[:, :, :2].reshape(-1, 2) - signal, dim=-1)
        loss += (distances[distances < safety_margin] ** 2).sum() * 5

    return loss


def calculate_curvature(trajectory):
    """计算轨迹的曲率"""
    # 计算速度
    velocities = trajectory[:, 1:, :2] - trajectory[:, :-1, :2]  # 计算相邻点之间的速度
    speeds = torch.norm(velocities, dim=-1)  # 计算速度的模

    # 计算加速度
    accelerations = velocities[:, 1:, :] - velocities[:, :-1, :]  # 计算相邻速度之间的加速度

    # 计算曲率
    cross_products = velocities[:, 1:, 0] * accelerations[:, :, 1] - velocities[:, 1:, 1] * accelerations[:, :, 0]
    curvature = torch.abs(cross_products) / (speeds[:,1:] ** 3 + 1e-6)  # 避免除以零

    return speeds, accelerations, curvature

def comfort_loss(pred_trajs):
    """
    计算舒适度损失
    :param pred_trajs: 预测的轨迹，形状为 (num_trajectories, num_points, 2)
    :return: 舒适度损失
    """
    c_loss = 0.0
    
    # 计算曲率和加速度
    speeds, accelerations, curvature = calculate_curvature(pred_trajs)
    
    # 曲率损失：曲率越大，损失越大
    c_loss += (curvature ** 2).sum()  # 曲率的平方和

    # 加速度损失：加速度变化越大，损失越大

    c_loss += (torch.norm(accelerations, dim=-1) ** 2).sum()  # 加速度的平方和

    return c_loss



def diversity_loss(pred_trajs, threshold=0.5):
    N, P, _ = pred_trajs.shape
    divers_loss = 0.0
    
    # 计算每对轨迹之间的平方距离
    # 使用广播机制计算距离矩阵
    diff = pred_trajs[:, None, :, :] - pred_trajs[None, :, :, :]
    squared_distances = torch.sum(diff ** 2, dim=-1)  # 计算平方距离

    # 仅保留上三角部分，避免自我比较
    squared_distances = squared_distances.triu(diagonal=1)

    # 计算损失
    mask = squared_distances < threshold ** 2  # 仅考虑小于阈值的平方距离
    divers_loss += torch.sum((threshold ** 2 - squared_distances[mask]))  # 计算损失

    return divers_loss


def distance_loss(pred_trajs):
    """轨迹长度"""
    # 获取轨迹数量和每条轨迹的点数
    N, P, _ = pred_trajs.shape
    
    # 计算相邻点之间的距离
    distances = torch.norm(pred_trajs[:, 1:, :2] -  pred_trajs[:, :-1, :2], dim=-1)  # 计算所有轨迹的相邻点之间的距离
    
    # 计算每条轨迹的总距离
    total_distance = distances.sum(dim=1)  # 对每条轨迹的距离求和
    
    # 计算总损失
    dis=_loss = total_distance.sum()  # 将所有轨迹的总距离相加

    return dis_loss


def add_random_noise(trajectory, noise_level=0.1):
    noise = torch.randn_like(trajectory) * noise_level  # 生成随机噪声
    return trajectory + noise  # 添加噪声


def calculate_average_distance(pred_trajs, gt_trajs):
    """计算所有预测轨迹与司机轨迹之间的平均距离"""
    # 计算每条预测轨迹与司机轨迹之间的距离
    distances = torch.norm(pred_trajs - gt_trajs.unsqueeze(0), dim=-1)  # 扩展维度以进行广播
    average_distances = distances.mean(dim=1)  # 计算每条轨迹的平均距离
    return average_distances

def local_similarity_loss(pred_trajs, gt_trajs,pred_scores,K=2):
    # 计算所有预测轨迹与司机轨迹之间的平均距离
    average_distances = calculate_average_distance(pred_trajs, gt_trajs)
    if pred_scores is not None:
        average_distances = average_distances * (K-pred_scores)
    # 找到最小距离及其对应的轨迹
    min_distance, min_index = torch.min(average_distances, dim=0)
    best_trajectory = pred_trajs[min_index]

    return min_distance.item(), best_trajectory

def calculate_occupancy_loss(pred_trajs, occupancy_map):
    """计算预测轨迹与占据地图的冲突损失"""
    
    # 获取轨迹数量和每条轨迹的点数
    N, P, _ = pred_trajs.shape
    
    # 将轨迹点的坐标转换为整数索引
    x = torch.floor(pred_trajs[:, :, 0]).long()  # 使用 floor 取整并转换为 long
    y = torch.floor(pred_trajs[:, :, 1]).long()
    z = torch.floor(pred_trajs[:, :, 2]).long()

    # 确保坐标在地图范围内
    x = torch.clamp(x, 0, occupancy_map.shape[0] - 1)
    y = torch.clamp(y, 0, occupancy_map.shape[1] - 1)
    z = torch.clamp(z, 0, occupancy_map.shape[2] - 1)

    # 计算被占据的点
    occupied_points = occupancy_map[x, y, z] == 1

    # 计算占据损失
    occupancy_loss = occupied_points.sum().item()  # 统计被占据的点的数量

    return occupancy_loss

def calculate_points_loss(pred_trajs, occupancy_map):
    """计算预测轨迹与点云冲突损失"""
    
    # 获取轨迹数量和每条轨迹的点数
    N, P, _ = pred_trajs.shape
    
    # 将轨迹点的坐标转换为整数索引
    x = torch.floor(pred_trajs[:, :, 0]).long()  # 使用 floor 取整并转换为 long
    y = torch.floor(pred_trajs[:, :, 1]).long()
    z = torch.floor(pred_trajs[:, :, 2]).long()

    # 确保坐标在地图范围内
    x = torch.clamp(x, 0, occupancy_map.shape[0] - 1)
    y = torch.clamp(y, 0, occupancy_map.shape[1] - 1)
    z = torch.clamp(z, 0, occupancy_map.shape[2] - 1)

    # 计算被占据的点
    occupied_points = occupancy_map[x, y, z] == 1

    # 计算占据损失
    occupancy_loss = occupied_points.sum().item()  # 统计被占据的点的数量

    return occupancy_loss
@LOSSES.register_module()
class motion_loss(nn.Module):
    def __init__(self):
        super(motion_loss, self).__init__()
        self.gt_trajs_with_scores=True
    def box_loss(self,pd_trajs,gt_boxes):
        return Bbox_loss(pd_trajs,gt_boxes)
    
    def map_loss(self,pd_trajs,gt_maps):
        pass
    def occ_loss(self,pd_trajs,gt_occ):
        pass
    def comfort_loss(self, pd_trajs):
        return comfort_loss(pd_trajs)
    
    def diversity_loss(self, pd_trajs):
        return diversity_loss(pd_trajs)
    
    def distance_loss(self,pd_trajs):
        return distance_loss(pd_trajs)
    
    def gt_loss(self,pd_trajs,gt_trajs,pd_scores=None):
        min_dist,idx = local_similarity_loss(pd_trajs,gt_trajs,pd_scores)
        return min_dist
    def forward(self, pd_traj, pd_scores, gt_traj, gt_boxes, gt_maps, gt_occ):
        """
            B:batch_size
            N: number of agents
            S: number of trajs
            T: number of time steps
            6: [x,y,cos_theta,sin_theta,vx,vy]
        pred_trajs: B x N x S x T x 6
        pred_scores: B x N x S
        gt_trajs: B x N x 1 x T x 6
        gt_boxes: B x N x 1 x 7 ->[x,y,z,w,l,h,yaw]
        gt_maps:
        gt_occ:
        """
        if gt_traj is not None:
            if self.gt_trajs_with_scores:
                gt_loss = self.gt_loss(pd_traj, gt_traj, pd_scores)
            else:
                gt_loss = self.gt_loss(pd_traj, gt_traj)
        else:
            gt_loss = 0
        if gt_boxes is not None:
            box_loss = self.box_loss(pd_traj, gt_boxes)
        else:
            box_loss = 0
        if gt_maps is not None:
            map_loss = self.map_loss(pd_traj, gt_maps)
        else:
            map_loss = 0
        if gt_occ is not None:
            occ_loss = self.occ_loss(pd_traj, gt_occ)
        else:
            occ_loss = 0
        comfort_loss = self.comfort_loss(pd_traj)
        diversity_loss = self.diversity_loss(pd_traj)
        distance_loss = self.distance_loss(pd_traj)
        loss = dict(
            gt_loss=gt_loss,
            box_loss=box_loss,
            map_loss=map_loss,
            occ_loss=occ_loss,
            comfort_loss=comfort_loss,
            diversity_loss=diversity_loss,
            distance_loss=distance_loss
        )
        return loss